package com.pai4j.aigc.llm;
import com.pai4j.aigc.config.properties.DeepseekLLMPropertiesConfig;
import com.pai4j.common.util.llm.DeepSeekAIClient;
import com.pai4j.domain.vo.llm.ChatCompletionMessage;
import com.pai4j.domain.vo.llm.ChatCompletionRequest;
import com.pai4j.domain.vo.llm.ChatCompletionResponse;
import com.pai4j.domain.vo.llm.ChatCompletionStreamChoice;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;

/**
 * @Author 程序员七七
 * @webSite https://www.coderutil.com
 * @Date 2024/10/29 23:15
 * @description
 */
@Service
public class DeepSeekAIService extends AbstractLLMChatService {

    @Autowired
    private DeepseekLLMPropertiesConfig deepseekLLMPropertiesConfig;

    /**
     * 同步对话
     *
     * @param messageList
     * @return
     * @throws Exception
     */
    public ChatCompletionResponse chat(List<ChatCompletionMessage> messageList) throws Exception {
        DeepSeekAIClient client = new DeepSeekAIClient(deepseekLLMPropertiesConfig.getApiKey());
        final List<ChatCompletionMessage> messages = messageList;
        return client.ChatCompletion(new ChatCompletionRequest(
                deepseekLLMPropertiesConfig.getModel(),
                messages,
                2000,
                0.3f,
                1
        ));
    }

    @Override
    protected ChatCompletionResponse chat(List<ChatCompletionMessage> messageList, Object config) throws Exception {
        if (config instanceof LLMRuntimeConfig runtime && runtime.getApiKey() != null) {
            String apiKey = runtime.getApiKey();
            String model = runtime.getModel() != null ? runtime.getModel() : deepseekLLMPropertiesConfig.getModel();
            DeepSeekAIClient client = new DeepSeekAIClient(apiKey);
            final List<ChatCompletionMessage> messages = messageList;
            Integer maxTokens = runtime.getMaxTokens() != null ? runtime.getMaxTokens() : 2000;
            Float temperature = runtime.getTemperature() != null ? runtime.getTemperature() : 0.3f;
            return client.ChatCompletion(new ChatCompletionRequest(
                    model,
                    messages,
                    maxTokens,
                    temperature,
                    1
            ));
        }
        return super.chat(messageList, config);
    }

    /**
     * 流式对话
     *
     * @param emitter
     * @param messageList
     * @throws Exception
     */
    public String chat(SseEmitter emitter, List<ChatCompletionMessage> messageList) throws Exception {

        // 用于记录流式推理完整结果，返回用于，用于对话消息持久化
        StringBuilder aiMessage = new StringBuilder();

        DeepSeekAIClient client = new DeepSeekAIClient(deepseekLLMPropertiesConfig.getApiKey());
        final List<ChatCompletionMessage> messages = messageList;
        try {
            client.ChatCompletionStream(new ChatCompletionRequest(
                    deepseekLLMPropertiesConfig.getModel(),
                    messages,
                    2000,
                    0.3f,
                    1
            )).subscribe(
                    streamResponse -> {
                        if (streamResponse.getChoices().isEmpty()) {
                            return;
                        }
                        for (ChatCompletionStreamChoice choice : streamResponse.getChoices()) {
                            String finishReason = choice.getFinishReason();
                            if (finishReason != null) {
                                emitter.send("finished");
                                continue;
                            }
                            String responseContent = choice.getDelta().getContent();
                            if (responseContent.equals("```") || responseContent.equals("json") ) {
                                continue;
                            }
                            responseContent = responseContent.replaceAll("\n", "<br>");
                            System.out.println(responseContent);
                            emitter.send(responseContent);
                            aiMessage.append(responseContent);
                        }
                    },
                    error -> {
                        error.printStackTrace();
                    },
                    () -> {
//                        emitter.complete(); // 完成事件流发送
                    }
            );
        } catch (Exception e) {
            e.printStackTrace();
        }
        return aiMessage.toString();
    }

    @Override
    protected String chat(SseEmitter emitter, List<ChatCompletionMessage> messageList, Object config) throws Exception {
        if (config instanceof LLMRuntimeConfig runtime && runtime.getApiKey() != null) {
            String apiKey = runtime.getApiKey();
            String model = runtime.getModel() != null ? runtime.getModel() : deepseekLLMPropertiesConfig.getModel();
            StringBuilder aiMessage = new StringBuilder();
            DeepSeekAIClient client = new DeepSeekAIClient(apiKey);
            final List<ChatCompletionMessage> messages = messageList;
            try {
                Integer maxTokens = runtime.getMaxTokens() != null ? runtime.getMaxTokens() : 2000;
                Float temperature = runtime.getTemperature() != null ? runtime.getTemperature() : 0.3f;
                client.ChatCompletionStream(new ChatCompletionRequest(
                        model,
                        messages,
                        maxTokens,
                        temperature,
                        1
                )).subscribe(
                        streamResponse -> {
                            if (streamResponse.getChoices().isEmpty()) {
                                return;
                            }
                            for (ChatCompletionStreamChoice choice : streamResponse.getChoices()) {
                                String finishReason = choice.getFinishReason();
                                if (finishReason != null) {
                                    emitter.send("finished");
                                    continue;
                                }
                                String responseContent = choice.getDelta().getContent();
                                if (responseContent.equals("```") || responseContent.equals("json") ) {
                                    return;
                                }
                                responseContent = responseContent.replaceAll("\n", "<br>");
                                emitter.send(responseContent);
                                aiMessage.append(responseContent);
                            }
                        },
                        error -> {
                            error.printStackTrace();
                        },
                        () -> {
                        }
                );
            } catch (Exception e) {
                e.printStackTrace();
            }
            return aiMessage.toString();
        }
        return super.chat(emitter, messageList, config);
    }
}